#!/bin/bash
## Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
# KubeOS is licensed under the Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#     http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
# PURPOSE.
## See the Mulan PSL v2 for more details.

set -e

ANOMALY=""
BENIGN=""
VALIDATE=false
MODELpath="../model"
EPOCH=50

function show_usage() {
   cat << EOF

Usage : sh auto-train [OPTIONS]

Options:
    -a              anomaly data file path (MUST in JSON format)
    -b              benign data file path for training (MUST in JSON format)
    -t              validate after training (with anomaly data provided)
    -m		    path to store models
    -e		    training epoches
    -h,--help       show help information
EOF
}

function param_parser() {

    while getopts "a:b:t:m:e:" opt
    do
      case $opt in
        a)
          ANOMALY="$OPTARG"
          ;;
	b)
          BENIGN="$OPTARG"
          ;;
	t)
          VALIDATE=$OPTARG
          ;;
	m)
          MODELpath="$OPTARG"
          ;;
	e)
          EPOCH=$OPTARG
          ;;
        *)
	  show_usage
	  exit 3
	  ;;
      esac
    done

}

function main() {
    set -e
    param_parser "$@"

    if [ -d "${MODELpath}" ]; then
        rm -rf ${MODELpath}/*
    else
	mkdir -p ${MODELpath}
    fi

    if [ -f "$BENIGN" ]; then
	cp "$BENIGN" ${MODELpath}/benign.json
	python3 process_behavior.py --d ${MODELpath} --file benign.json
    else
	echo "benign data [$BENIGN] NOT exists."
	exit 2
    fi

    if [ -f "$ANOMALY" ]; then
        cp "$ANOMALY" ${MODELpath}/anomaly.json
        python3 process_behavior.py --d ${MODELpath} --file anomaly.json
    else
        echo "anomaly data [$ANOMALY] NOT exists."
    fi

    python3 filename-embedding.py --d ${MODELpath} 
    python3 cmdline-embedding.py --d ${MODELpath} 
    python3 caculate-weight.py --d ${MODELpath} 
    if [ -f  "${MODELpath}/process-event-anomaly.txt" ]; then
	python3 train.py --epoch ${EPOCH} --d ${MODELpath} --v True > ${MODELpath}/train.log
    else
	python3 train.py --epoch ${EPOCH} --d ${MODELpath} > ${MODELpath}/train.log
    fi

    set +e
    ret=$(cat ${MODELpath}/train.log | grep -o "Bad Training")
    if [ "$ret" == "" ]; then
	cat ${MODELpath}/threshold.txt
	echo "Complete Training"
    else
	echo "***[Bad Training] -- increase epoches or modify min_epoch_loss ***"
    fi

}

main "$@"

